package com.webank.maling.documentation.repository.graph;

import com.webank.maling.base.model.MethodInfo;
import com.webank.maling.base.model.SubgraphData;
import com.webank.maling.documentation.repository.graph.support.GraphSupport;
import com.webank.maling.repository.nebula.NebulaGraphClient;
import com.vesoft.nebula.client.graph.data.Node;
import com.vesoft.nebula.client.graph.data.Relationship;
import com.vesoft.nebula.client.graph.data.ResultSet;
import com.vesoft.nebula.client.graph.data.ValueWrapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 子图查询服务
 * 负责从NebulaGraph中查询和解析子图数据
 *
 * @author diodehe
 */
@Slf4j
@Service
public class SubgraphRepository {

    private final NebulaGraphClient nebulaGraphClient;

    @Autowired
    public SubgraphRepository(NebulaGraphClient nebulaGraphClient) {
        this.nebulaGraphClient = nebulaGraphClient;
    }
    
    /**
     * 根据层级获取子图数据
     * 
     * @param entryPointId 入口点ID
     * @param maxSteps 最大步数
     * @return 子图数据
     */
    public SubgraphData getSubgraphByLevel(String entryPointId, int maxSteps) {
        try {
            log.info("开始查询入口点 {} 的 {} 步子图", entryPointId, maxSteps);
            
            // 构建子图查询语句
            String query = """
                    GET SUBGRAPH WITH PROP %d STEPS FROM "%s" OUT calls,out_calls,implemented_by,overridden_by,super_calls,interface_calls,subtype_calls,injection_calls
                    YIELD VERTICES AS nodes, EDGES AS relationships
                    """.formatted(maxSteps, entryPointId);
            
            ResultSet resultSet = nebulaGraphClient.execute(query);
            
            if (!resultSet.isSucceeded()) {
                log.error("子图查询失败: {}", resultSet.getErrorMessage());
                return createEmptySubgraph();
            }
            
            // 解析查询结果
            SubgraphData subgraph = parseSubgraphResult(resultSet, entryPointId);
            
            log.info("成功查询到子图，节点数: {}, 边数: {}, 最大层级: {}", 
                    subgraph.getTotalNodes(), subgraph.getTotalEdges(), subgraph.getMaxLevel());
            
            return subgraph;
            
        } catch (Exception e) {
            log.error("查询子图时发生错误", e);
            return createEmptySubgraph();
        }
    }
    
    /**
     * 获取指定步数的子图
     * 
     * @param entryPointId 入口点ID
     * @param steps 步数
     * @return 子图数据
     */
    public SubgraphData getSubgraph(String entryPointId, int steps) {
        return getSubgraphByLevel(entryPointId, steps);
    }
    
    /**
     * 解析子图查询结果
     */
    public SubgraphData parseSubgraphResult(ResultSet resultSet, String entryPointId) {
        try {
            if (resultSet.rowsSize() == 0) {
                log.warn("子图查询结果为空");
                return createEmptySubgraph();
            }

            // 初始化合并后的数据容器
            Map<String, MethodInfo> allMethods = new HashMap<>();
            List<SubgraphData.CallRelation> allRelations = new ArrayList<>();

            // 循环处理所有行的数据
            for (int i = 0; i < resultSet.rowsSize(); i++) {
                try {
                    log.debug("处理第 {} 行子图数据", i + 1);

                    // 获取当前行的节点和边数据
                    ValueWrapper nodesWrapper = resultSet.rowValues(i).get(0); // nodes
                    ValueWrapper edgesWrapper = resultSet.rowValues(i).get(1);  // relationships

                    // 解析并合并节点数据
                    Map<String, MethodInfo> rowMethods = parseNodes(nodesWrapper);
                    allMethods.putAll(rowMethods);
                    log.debug("第 {} 行解析到 {} 个节点", i + 1, rowMethods.size());

                    // 解析并合并边数据
                    List<SubgraphData.CallRelation> rowRelations = parseEdges(edgesWrapper);
                    allRelations.addAll(rowRelations);
                    log.debug("第 {} 行解析到 {} 条边", i + 1, rowRelations.size());

                } catch (Exception e) {
                    log.warn("处理第 {} 行子图数据时发生错误，跳过该行: {}", i + 1, e.getMessage());
                }
            }

            log.info("合并后共解析到 {} 个节点，{} 条边", allMethods.size(), allRelations.size());

            // 检查是否有有效数据
            if (allMethods.isEmpty()) {
                log.warn("所有行处理完成后未获取到任何节点数据");
                return createEmptySubgraph();
            }

            // 计算每个节点的层级
            Map<Integer, List<MethodInfo>> methodsByLevel = calculateNodeLevels(allMethods, allRelations, entryPointId);

            // 获取入口点信息
            MethodInfo entryPoint = allMethods.get(entryPointId);
            if (entryPoint == null) {
                log.warn("在子图中未找到入口点: {}", entryPointId);
                return createEmptySubgraph();
            }

            // 将所有方法转换为列表
            List<MethodInfo> allMethodsList = methodsByLevel.values().stream()
                    .flatMap(List::stream)
                    .collect(Collectors.toList());

            // 构建子图数据
            return SubgraphData.builder()
                    .entryPoint(entryPoint)
                    .methods(allMethodsList)
                    .relations(allRelations)
                    .maxLevel(methodsByLevel.keySet().stream().max(Integer::compareTo).orElse(0))
                    .timestamp(System.currentTimeMillis())
                    .build();

        } catch (Exception e) {
            log.error("解析子图结果时发生错误", e);
            return createEmptySubgraph();
        }
    }
    
    /**
     * 解析节点数据
     */
    private Map<String, MethodInfo> parseNodes(ValueWrapper nodesWrapper) {
        Map<String, MethodInfo> methods = new HashMap<>();
        
        try {
            if (!nodesWrapper.isList()) {
                return methods;
            }
            
            var nodesList = nodesWrapper.asList();
            
            for (ValueWrapper nodeWrapper : nodesList) {
                if (nodeWrapper.isVertex()) {
                    Node node = nodeWrapper.asNode();
                    String nodeId = node.getId().asString();
                    
                    // 只处理function类型的节点
                    if (node.tagNames().contains("function")) {
                        MethodInfo methodInfo = GraphSupport.parseVertexToMethodInfo(node);
                        if (methodInfo != null) {
                            methods.put(nodeId, methodInfo);
                        }
                    }
                }
            }
            
        } catch (Exception e) {
            log.warn("解析节点数据时发生错误", e);
        }
        
        return methods;
    }
    
    /**
     * 解析边数据
     */
    private List<SubgraphData.CallRelation> parseEdges(ValueWrapper edgesWrapper) {
        List<SubgraphData.CallRelation> relations = new ArrayList<>();
        
        try {
            if (!edgesWrapper.isList()) {
                return relations;
            }
            
            var edgesList = edgesWrapper.asList();
            
            for (ValueWrapper edgeWrapper : edgesList) {
                if (edgeWrapper.isEdge()) {
                    Relationship relationship = edgeWrapper.asRelationship();

                    SubgraphData.CallRelation relation = SubgraphData.CallRelation.builder()
                            .sourceMethodId(relationship.srcId().asString())
                            .targetMethodId(relationship.dstId().asString())
                            .relationType(relationship.edgeName())
                            .lineNumber(GraphSupport.getEdgeIntProperty(relationship, "line_number"))
                            .dependencyType(GraphSupport.getEdgeStringProperty(relationship, "dependency_type"))
                            .build();
                    
                    relations.add(relation);
                }
            }
            
        } catch (Exception e) {
            log.warn("解析边数据时发生错误", e);
        }
        
        return relations;
    }
    
    /**
     * 计算节点层级（基于BFS算法）
     */
    private Map<Integer, List<MethodInfo>> calculateNodeLevels(Map<String, MethodInfo> allMethods, 
                                                              List<SubgraphData.CallRelation> relations, 
                                                              String entryPointId) {
        
        Map<Integer, List<MethodInfo>> methodsByLevel = new HashMap<>();
        Map<String, Integer> nodeLevels = new HashMap<>();
        
        // 构建邻接表
        Map<String, List<String>> adjacencyList = new HashMap<>();
        for (SubgraphData.CallRelation relation : relations) {
            adjacencyList.computeIfAbsent(relation.getSourceMethodId(), k -> new ArrayList<>())
                    .add(relation.getTargetMethodId());
        }
        
        // BFS计算层级
        Queue<String> queue = new LinkedList<>();
        Set<String> visited = new HashSet<>();
        
        // 入口点为第0层
        queue.offer(entryPointId);
        visited.add(entryPointId);
        nodeLevels.put(entryPointId, 0);
        
        while (!queue.isEmpty()) {
            String currentNode = queue.poll();
            int currentLevel = nodeLevels.get(currentNode);
            
            List<String> neighbors = adjacencyList.getOrDefault(currentNode, List.of());
            for (String neighbor : neighbors) {
                if (!visited.contains(neighbor) && allMethods.containsKey(neighbor)) {
                    visited.add(neighbor);
                    int neighborLevel = currentLevel + 1;
                    nodeLevels.put(neighbor, neighborLevel);
                    queue.offer(neighbor);
                }
            }
        }
        
        // 按层级分组
        for (Map.Entry<String, Integer> entry : nodeLevels.entrySet()) {
            String nodeId = entry.getKey();
            Integer level = entry.getValue();
            
            MethodInfo method = allMethods.get(nodeId);
            if (method != null) {
                method.setLevel(level);
                methodsByLevel.computeIfAbsent(level, k -> new ArrayList<>()).add(method);
            }
        }
        
        return methodsByLevel;
    }
    
    
    /**
     * 创建空子图
     */
    private SubgraphData createEmptySubgraph() {
        return SubgraphData.builder()
                .methods(new ArrayList<>())
                .relations(new ArrayList<>())
                .maxLevel(0)
                .timestamp(System.currentTimeMillis())
                .build();
    }
}
